import hashlib
from typing import Dict, Any, Optional, Union
import json
import time
from kg_retrieval.env_config import REDIS_CONFIG

class CacheManager:
    def __init__(self, host=None, port=None, db=None, ttl=3600):
        # 使用环境变量或参数
        self.host = host or REDIS_CONFIG["REDIS_HOST"]
        self.port = port or REDIS_CONFIG["REDIS_PORT"]
        self.db = db if db is not None else REDIS_CONFIG["REDIS_DB"]
        self.redis_enabled = REDIS_CONFIG["REDIS_ENABLED"]
        self.redis_password = REDIS_CONFIG["REDIS_PASSWORD"]
        
        try:
            if self.redis_enabled:
                import redis
                self.redis = redis.Redis(
                    host=self.host, 
                    port=self.port, 
                    db=self.db,
                    password=self.redis_password if self.redis_password else None
                )
                # 测试连接
                self.redis.ping()
                print(f"Redis连接成功: {self.host}:{self.port}")
                self.mode = "redis"
            else:
                print("Redis缓存未启用，使用内存缓存")
                self.redis = None
                self.memory_cache = {}  # 简单的内存缓存
                self.memory_expiry = {}  # 键的过期时间
                self.mode = "memory"
        except (ImportError, Exception) as e:
            print(f"Redis连接失败，将使用内存缓存: {e}")
            self.redis = None
            self.memory_cache = {}  # 简单的内存缓存
            self.memory_expiry = {}  # 键的过期时间
            self.mode = "memory"
        
        self.ttl = ttl  # 默认缓存时间1小时
    
    def get_cache_key(self, query: str, company_id: Optional[str] = None, use_overlap_segments: bool = False, full_content_merge_neigh: bool = True) -> str:
        """生成缓存键，包括重叠分段和内容合并设置"""
        if company_id:
            key_data = f"{query}:{company_id}:{use_overlap_segments}:{full_content_merge_neigh}"
        else:
            key_data = f"{query}:{use_overlap_segments}:{full_content_merge_neigh}"
        
        return f"search:{hashlib.md5(key_data.encode()).hexdigest()}"
    
    def get(self, key: str) -> Optional[Dict[str, Any]]:
        """获取缓存数据"""
        if self.mode == "redis" and self.redis:
            data = self.redis.get(key)
            if data:
                try:
                    return json.loads(data)
                except:
                    pass
        elif self.mode == "memory":
            # 检查是否过期
            if key in self.memory_expiry and time.time() <= self.memory_expiry[key]:
                return self.memory_cache.get(key)
            elif key in self.memory_cache:
                # 如果过期，删除缓存
                del self.memory_cache[key]
                if key in self.memory_expiry:
                    del self.memory_expiry[key]
        
        return None
    
    def set(self, key: str, value: Union[Dict[str, Any], list, str], ttl: Optional[int] = None) -> bool:
        """设置缓存数据"""
        ttl = ttl or self.ttl
        
        try:
            # 将数据序列化为JSON
            json_data = json.dumps(value)
            
            if self.mode == "redis" and self.redis:
                self.redis.setex(key, ttl, json_data)
            else:
                self.memory_cache[key] = value
                self.memory_expiry[key] = time.time() + ttl
            
            return True
        except Exception as e:
            print(f"缓存设置失败: {e}")
            return False
    
    def clear_cache(self, pattern: str = "search:*") -> int:
        """清除缓存"""
        if self.mode == "redis" and self.redis:
            keys = self.redis.keys(pattern)
            if keys:
                return self.redis.delete(*keys)
            return 0
        else:
            # 清除内存缓存中的匹配项
            # 这里简单处理，清除所有以pattern开头的键（去掉通配符*）
            prefix = pattern.rstrip("*")
            keys_to_delete = [k for k in self.memory_cache.keys() if k.startswith(prefix)]
            
            for k in keys_to_delete:
                del self.memory_cache[k]
                if k in self.memory_expiry:
                    del self.memory_expiry[k]
            
            return len(keys_to_delete) 